# -*- coding: utf-8 -*-
import time
import base64
import json
import requests
import glob
import os
import shutil
from PIL import Image
import cv2
from io import BytesIO

yolo_url = "http://yj-ocr-all-wan.haofenshu.com/analyse_layout"


class YoloAPI:
    """
    YOLO检测服务的接口类
    """
    def __init__(self, yolo_url="http://yj-ocr-all-wan.haofenshu.com/analyse_layout"):
        self.yolo_url = yolo_url

    def analyse_layout_of_cv_image(self, cv_image, select_types=None):
        '''
        调yolo检测的api获取检测的结果，如果select_types为None，则返回所有的结果，如果select_types不为None，则只返回列表中的类型，如
        select_types = ['figure'] 则只返回插图的类型
        '''

        image_base64 = self.cv_image_data_to_base64(cv_image)
        if image_base64 is None:
            return []

        data = {
            "image": image_base64,
        }
        st = time.time()
        response = requests.post(self.yolo_url, data=json.dumps(data))  # 修改为你的API地址

        print("yolo api status : ", response.status_code)
        if response.status_code == 200:
            #print(response.text)
            response_data = json.loads(response.text)
            layout_boxes = response_data['data']
            if layout_boxes and select_types:
                # 要根据提供的类型做筛选
                selected_boxes = []
                for box in layout_boxes:
                    box_name = box["name"]
                    if box_name in select_types:
                        selected_boxes.append(box)
                return selected_boxes
        else:
            layout_boxes = []
        return layout_boxes

    def image_path_to_base64(self, image_path):
        with open(image_path, "rb") as image_file:
            return base64.b64encode(image_file.read()).decode('utf-8')
        

    def cv_image_data_to_base64(self, image, img_format='PNG'):
        '''
        cv2格式图片转base64
        '''
        try:
            if image is None:
                return None
            if len(image.shape) == 2:
                pil_image = Image.fromarray(image)
            else:
                pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
            img_buffer = BytesIO()
            pil_image.save(img_buffer, format=img_format)
            byte_data = img_buffer.getvalue()
            base64_str = base64.b64encode(byte_data)
            base64_str = str(base64_str, "utf-8")  # bytes -> base64 string
            return base64_str
        except Exception as e:
            print('img2base64 failed.')
            return None


    def visual_result(self, file_path, layout_boxes, output_path):

        colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0), (0, 255, 255), (255, 0, 255), (255, 165, 0), (255, 192, 203), (165, 42, 42),  (0, 0, 139), (0, 100, 0), (0, 139, 139), (0, 255, 0)]
        types = ["title", "plain_text", "abandon", "figure", "figure_caption", "table", "table_caption", "table_footnote", "isolate_formula", "formula_caption", "embedding", "isolated", "text"]


        image = cv2.imread(file_path, -1)

        c = 1
        if len(image.shape) > 2:
            h, w, c = image.shape
        if c == 1:
            image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)

        type_counts = dict()
        for ii, layout_box in enumerate(layout_boxes):

            type_of_box = layout_box["name"]
            class_id = layout_box["class"]
            coordinate_of_box = layout_box["box"]
            x1, y1, x2, y2 = coordinate_of_box["x1"], coordinate_of_box["y1"], coordinate_of_box["x2"], coordinate_of_box["y2"]

            if type_of_box in type_counts:
                type_counts[type_of_box] = type_counts[type_of_box] +1
            else:
                type_counts[type_of_box] = 1

            cv2.rectangle(image, (int(x1), int(y1)), (int(x2), int(y2)), colors[class_id], 2)
            cv2.putText(image, type_of_box, (int(x1 + 20), int(y1) + 20), cv2.FONT_HERSHEY_SIMPLEX, 0.5, colors[class_id], 2)

        # output_path = os.path.join("../output", "structure_" + os.path.basename(file_path))
        if output_path is not None:
            cv2.imwrite(output_path, image)





    

if __name__ == "__main__":

    files = glob.glob("./data/*.jpg") + glob.glob("./data/*.png") + glob.glob("./data/*.JPG")

    files = sorted(files)

    yolo_api = YoloAPI(yolo_url)

    for file in files:

        basename = os.path.basename(file)
        output_folder = "./output"
        if not os.path.exists(output_folder):
            os.mkdir(output_folder)
        output_path = os.path.join(output_folder, basename)

        image = cv2.imread(file, 0)
        layout_boxes = yolo_api.analyse_layout_of_cv_image(image, select_types=['figure'])
        print("layout_boxes : ", layout_boxes)
        if layout_boxes is not None:
            yolo_api.visual_result(file, layout_boxes, output_path)
        else:
            print("layout_boxes is none")



        

